# 向量检索（Chroma 替换版）

import chromadb
from chromadb.config import Settings
import os
from langchain_community.embeddings.dashscope import DashScopeEmbeddings

# 设置环境变量
os.environ["DASHSCOPE_API_KEY"] = "sk-e9920e5c211248f28d4b4750e8eea906"
import logging

# 配置日志记录
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
# self.embeddings = DashScopeEmbeddings(
#             model='text-embedding-v1',
#             dashscope_api_key=os.getenv("DASHSCOPE_API_KEY")
#         )
def split_text(text, max_length=2048):      # 文本分割函数
    return [text[i:i+max_length] for i in range(0, len(text), max_length)]

class ChromaRetriever:
    """Chroma向量检索工具类，专注于向量检索和职位分区管理（用metadata实现分区）"""

    def __init__(self):
        try:
            self.client = chromadb.Client(Settings(persist_directory="./chroma_db"))
            self.collection = self.client.get_or_create_collection("resume_vectors")
            self.embeddings = DashScopeEmbeddings(
                model='text-embedding-v1'
            )
        except Exception as e:
            logger.error(f"初始化 DashScopeEmbeddings 失败: {e}")
            raise

    def sync_vector(self, jobposting_id, resume_id, text):
        # 分段   长文本分段嵌入 + 向量合并
        text_chunks = split_text(text)
        if not text_chunks:
            text_chunks = [""]  # 避免空文本

        # 逐个分段嵌入
        embeddings = []
        for chunk in text_chunks:
            vector = self.embeddings.embed_query(chunk)
            embeddings.append(vector)

        # 合并向量（取平均或其他策略）
        merged_vector = [sum(col) / len(col) for col in zip(*embeddings)]

        self.collection.add(
            ids=[str(resume_id)],
            metadatas=[{"jobposting_id": jobposting_id, "resume_id": str(resume_id)}],
            documents=[text],  # 存储完整文本
            embeddings=[merged_vector]
        )

    def search(self, jobposting_id, query_text, top_k=20):
        """在指定职位分区内执行向量检索（用metadata过滤）"""
        query_vector = self.embeddings.embed_query(query_text)
        results = self.collection.query(
            query_embeddings=[query_vector],
            n_results=top_k,
            where={"jobposting_id": jobposting_id},
            include=["metadatas", "distances"]  # 确保返回元数据
        )
        # 从元数据提取业务 resume_id，关联向量库 ID 和距离
        chroma_scores = {}
        for res_id, dist, meta in zip(
                results["ids"][0],
                results["distances"][0],
                results["metadatas"][0]
        ):
            business_resume_id = meta["resume_id"]  # 元数据中存的业务 ID
            chroma_scores[business_resume_id] = dist  # 用业务 ID 作为键
        print("chroma：",chroma_scores)
        return chroma_scores
